import {
    ArrayExpr,
    DataModel,
    DataModelAttribute,
    DataModelField,
    Expression,
    isArrayExpr,
    isBooleanLiteral,
    isDataModel,
    isDataModelField,
    isInvocationExpr,
    isNumberLiteral,
    isObjectExpr,
    isReferenceExpr,
    isStringLiteral,
    isTypeDef,
    ObjectExpr,
    ReferenceExpr,
    TypeDef,
    TypeDefField,
} from '@zenstackhq/language/ast';
import type { RuntimeAttribute } from '@zenstackhq/runtime';
import { lowerCaseFirst } from '@zenstackhq/runtime/local-helpers';
import { streamAst } from 'langium';
import { FunctionDeclarationStructure, OptionalKind, Project, VariableDeclarationKind } from 'ts-morph';
import { match } from 'ts-pattern';
import {
    CodeWriter,
    ExpressionContext,
    FastWriter,
    getAttribute,
    getAttributeArg,
    getAttributeArgs,
    getAuthDecl,
    getDataModels,
    getInheritedFromDelegate,
    getLiteral,
    getRelationBackLink,
    getRelationField,
    hasAttribute,
    isAuthInvocation,
    isEnumFieldReference,
    isForeignKeyField,
    isIdField,
    resolved,
    saveSourceFile,
    TypeScriptExpressionTransformer,
} from '.';

/**
 * Options for generating model metadata
 */
export type ModelMetaGeneratorOptions = {
    /**
     * Output directory
     */
    output: string;

    /**
     * Whether to generate all attributes
     */
    generateAttributes: boolean;

    /**
     * Whether to preserve the pre-compilation TypeScript files
     */
    preserveTsFiles?: boolean;

    /**
     * Map from full names to shortened names, used for extra fields/relations generated by ZenStack
     */
    shortNameMap?: Map<string, string>;
};

export function generate(
    project: Project,
    models: DataModel[],
    typeDefs: TypeDef[],
    options: ModelMetaGeneratorOptions
) {
    const sf = project.createSourceFile(options.output, undefined, { overwrite: true });

    // generate: import type { ModelMeta } from '@zenstackhq/runtime';
    sf.addImportDeclaration({
        isTypeOnly: true,
        namedImports: ['ModelMeta'],
        moduleSpecifier: '@zenstackhq/runtime',
    });

    const writer = new FastWriter();
    const extraFunctions: OptionalKind<FunctionDeclarationStructure>[] = [];
    generateModelMetadata(models, typeDefs, writer, options, extraFunctions);

    sf.addVariableStatement({
        declarationKind: VariableDeclarationKind.Const,
        declarations: [{ name: 'metadata', type: 'ModelMeta', initializer: writer.result }],
    });

    if (extraFunctions.length > 0) {
        sf.addFunctions(extraFunctions);
    }

    sf.addStatements('export default metadata;');

    if (options.preserveTsFiles) {
        saveSourceFile(sf);
    }

    return sf;
}

function generateModelMetadata(
    dataModels: DataModel[],
    typeDefs: TypeDef[],
    writer: CodeWriter,
    options: ModelMetaGeneratorOptions,
    extraFunctions: OptionalKind<FunctionDeclarationStructure>[]
) {
    writer.block(() => {
        writeModels(writer, dataModels, options, extraFunctions);
        writeTypeDefs(writer, typeDefs, options, extraFunctions);
        writeDeleteCascade(writer, dataModels);
        writeShortNameMap(options, writer);
        writeAuthModel(writer, dataModels, typeDefs);
    });
}

function writeModels(
    writer: CodeWriter,
    dataModels: DataModel[],
    options: ModelMetaGeneratorOptions,
    extraFunctions: OptionalKind<FunctionDeclarationStructure>[]
) {
    writer.write('models:');
    writer.block(() => {
        for (const model of dataModels) {
            writer.write(`${lowerCaseFirst(model.name)}:`);
            writer.block(() => {
                writer.write(`name: '${model.name}',`);
                writeBaseTypes(writer, model);
                writeFields(writer, model, options, extraFunctions);
                writeUniqueConstraints(writer, model);
                if (options.generateAttributes) {
                    writeModelAttributes(writer, model);
                }
                writeDiscriminator(writer, model);
            });
            writer.writeLine(',');
        }
    });
    writer.writeLine(',');
}

function writeTypeDefs(
    writer: CodeWriter,
    typedDefs: TypeDef[],
    options: ModelMetaGeneratorOptions,
    extraFunctions: OptionalKind<FunctionDeclarationStructure>[]
) {
    if (typedDefs.length === 0) {
        return;
    }
    writer.write('typeDefs:');
    writer.block(() => {
        for (const typeDef of typedDefs) {
            writer.write(`${lowerCaseFirst(typeDef.name)}:`);
            writer.block(() => {
                writer.write(`name: '${typeDef.name}',`);
                writeFields(writer, typeDef, options, extraFunctions);
            });
            writer.writeLine(',');
        }
    });
    writer.writeLine(',');
}

function writeBaseTypes(writer: CodeWriter, model: DataModel) {
    if (model.superTypes.length > 0) {
        writer.write('baseTypes: [');
        writer.write(model.superTypes.map((t) => `'${t.ref?.name}'`).join(', '));
        writer.write('],');
    }
}

function writeAuthModel(writer: CodeWriter, dataModels: DataModel[], typeDefs: TypeDef[]) {
    const authModel = getAuthDecl([...dataModels, ...typeDefs]);
    if (authModel) {
        writer.writeLine(`authModel: '${authModel.name}'`);
    }
}

function writeDeleteCascade(writer: CodeWriter, dataModels: DataModel[]) {
    writer.write('deleteCascade:');
    writer.block(() => {
        for (const model of dataModels) {
            const cascades = getDeleteCascades(model);
            if (cascades.length > 0) {
                writer.writeLine(`${lowerCaseFirst(model.name)}: [${cascades.map((n) => `'${n}'`).join(', ')}],`);
            }
        }
    });
    writer.writeLine(',');
}

function writeUniqueConstraints(writer: CodeWriter, model: DataModel) {
    const constraints = getUniqueConstraints(model);
    if (constraints.length > 0) {
        writer.write('uniqueConstraints:');
        writer.block(() => {
            for (const constraint of constraints) {
                writer.write(`${constraint.name}: {
                                name: "${constraint.name}",
                                fields: ${JSON.stringify(constraint.fields)}
                            },`);
            }
        });
        writer.write(',');
    }
}

function writeModelAttributes(writer: CodeWriter, model: DataModel) {
    const attrs = getAttributes(model);
    if (attrs.length > 0) {
        writer.write(`
attributes: ${JSON.stringify(attrs)},`);
    }
}

function writeDiscriminator(writer: CodeWriter, model: DataModel) {
    const delegateAttr = getAttribute(model, '@@delegate');
    if (!delegateAttr) {
        return;
    }
    const discriminator = getAttributeArg(delegateAttr, 'discriminator') as ReferenceExpr;
    if (!discriminator) {
        return;
    }
    if (discriminator) {
        writer.write(`discriminator: ${JSON.stringify(discriminator.target.$refText)},`);
    }
}

function writeFields(
    writer: CodeWriter,
    container: DataModel | TypeDef,
    options: ModelMetaGeneratorOptions,
    extraFunctions: OptionalKind<FunctionDeclarationStructure>[]
) {
    writer.write('fields:');
    writer.block(() => {
        for (const f of container.fields) {
            const dmField = isDataModelField(f) ? f : undefined;

            writer.write(`${f.name}: {`);

            writer.write(`
        name: "${f.name}",
        type: "${
            f.type.reference
                ? f.type.reference.$refText
                : // eslint-disable-next-line @typescript-eslint/no-non-null-assertion
                  f.type.type!
        }",`);

            if (dmField && isIdField(dmField)) {
                writer.write(`
        isId: true,`);
            }

            if (isDataModel(f.type.reference?.ref)) {
                writer.write(`
        isDataModel: true,`);
            } else if (isTypeDef(f.type.reference?.ref)) {
                writer.write(`
        isTypeDef: true,`);
            }

            if (f.type.array) {
                writer.write(`
        isArray: true,`);
            }

            if (f.type.optional) {
                writer.write(`
        isOptional: true,`);
            }

            if (options.generateAttributes) {
                const attrs = getAttributes(f);
                if (attrs.length > 0) {
                    writer.write(`
        attributes: ${JSON.stringify(attrs)},`);
                }
            } else {
                // only include essential attributes
                const attrs = getAttributes(f).filter((attr) => ['@default', '@updatedAt'].includes(attr.name));
                if (attrs.length > 0) {
                    writer.write(`
        attributes: ${JSON.stringify(attrs)},`);
                }
            }

            const defaultValueProvider = generateDefaultValueProvider(f, extraFunctions);
            if (defaultValueProvider) {
                writer.write(`
            defaultValueProvider: ${defaultValueProvider},`);
            }

            if (dmField) {
                // metadata specific to DataModelField

                const backlink = getRelationBackLink(dmField);
                const fkMapping = generateForeignKeyMapping(dmField);

                if (backlink) {
                    writer.write(`
        backLink: '${backlink.name}',`);
                }

                if (isRelationOwner(dmField, backlink)) {
                    writer.write(`
        isRelationOwner: true,`);
                }

                const onDeleteAction = getOnDeleteAction(dmField);
                if (onDeleteAction) {
                    writer.write(`
        onDeleteAction: '${onDeleteAction}',`);
                }

                const onUpdateAction = getOnUpdateAction(dmField);
                if (onUpdateAction) {
                    writer.write(`
        onUpdateAction: '${onUpdateAction}',`);
                }

                if (isForeignKeyField(dmField)) {
                    writer.write(`
        isForeignKey: true,`);
                    const relationField = getRelationField(dmField);
                    if (relationField) {
                        writer.write(`
        relationField: '${relationField.name}',`);
                    }
                }

                if (fkMapping && Object.keys(fkMapping).length > 0) {
                    writer.write(`
        foreignKeyMapping: ${JSON.stringify(fkMapping)},`);
                }

                const inheritedFromDelegate = getInheritedFromDelegate(dmField);
                if (inheritedFromDelegate && !isIdField(dmField)) {
                    writer.write(`
        inheritedFrom: ${JSON.stringify(inheritedFromDelegate.name)},`);
                }

                if (isAutoIncrement(dmField)) {
                    writer.write(`
        isAutoIncrement: true,`);
                }
            }

            writer.write(`
    },`);
        }
    });
    writer.write(',');
}

function getAttributes(target: DataModelField | DataModel | TypeDefField): RuntimeAttribute[] {
    return target.attributes
        .map((attr) => {
            const args: Array<{ name?: string; value?: unknown }> = [];
            for (const arg of attr.args) {
                const argName = arg.$resolvedParam?.name ?? arg.name;
                const argValue = exprToValue(arg.value);
                args.push({ name: argName, value: argValue });
            }
            return { name: resolved(attr.decl).name, args };
        })
        .filter((d): d is RuntimeAttribute => !!d);
}

function getUniqueConstraints(model: DataModel) {
    const constraints: Array<{ name: string; fields: string[] }> = [];

    const extractConstraint = (attr: DataModelAttribute) => {
        const argsMap = getAttributeArgs(attr);
        if (argsMap.fields) {
            const fieldNames = (argsMap.fields as ArrayExpr).items.map(
                (item) => resolved((item as ReferenceExpr).target).name
            );
            let constraintName = argsMap.name && getLiteral<string>(argsMap.name);
            if (!constraintName) {
                // default constraint name is fields concatenated with underscores
                constraintName = fieldNames.join('_');
            }
            return { name: constraintName, fields: fieldNames };
        } else {
            return undefined;
        }
    };

    const addConstraint = (constraint: { name: string; fields: string[] }) => {
        if (!constraints.some((c) => c.name === constraint.name)) {
            constraints.push(constraint);
        }
    };

    // field-level @id first
    for (const field of model.fields) {
        if (hasAttribute(field, '@id')) {
            addConstraint({ name: field.name, fields: [field.name] });
        }
    }

    // then model-level @@id
    for (const attr of model.attributes.filter((attr) => attr.decl.ref?.name === '@@id')) {
        const constraint = extractConstraint(attr);
        if (constraint) {
            addConstraint(constraint);
        }
    }

    // then field-level @unique
    for (const field of model.fields) {
        if (hasAttribute(field, '@unique')) {
            addConstraint({ name: field.name, fields: [field.name] });
        }
    }

    // then model-level @@unique
    for (const attr of model.attributes.filter((attr) => attr.decl.ref?.name === '@@unique')) {
        const constraint = extractConstraint(attr);
        if (constraint) {
            addConstraint(constraint);
        }
    }

    return constraints;
}

function isRelationOwner(field: DataModelField, backLink: DataModelField | undefined) {
    if (!isDataModel(field.type.reference?.ref)) {
        return false;
    }

    if (!backLink) {
        // CHECKME: can this really happen?
        return true;
    }

    if (!hasAttribute(field, '@relation') && !hasAttribute(backLink, '@relation')) {
        // if neither side has `@relation` attribute, it's an implicit many-to-many relation,
        // both sides are owners
        return true;
    }

    return holdsForeignKey(field);
}

function holdsForeignKey(field: DataModelField) {
    const relation = field.attributes.find((attr) => attr.decl.ref?.name === '@relation');
    if (!relation) {
        return false;
    }
    const fields = getAttributeArg(relation, 'fields');
    return !!fields;
}

function generateForeignKeyMapping(field: DataModelField) {
    const relation = field.attributes.find((attr) => attr.decl.ref?.name === '@relation');
    if (!relation) {
        return undefined;
    }
    const fields = getAttributeArg(relation, 'fields');
    const references = getAttributeArg(relation, 'references');
    if (!isArrayExpr(fields) || !isArrayExpr(references) || fields.items.length !== references.items.length) {
        return undefined;
    }

    const fieldNames = fields.items.map((item) => (isReferenceExpr(item) ? item.target.$refText : undefined));
    const referenceNames = references.items.map((item) => (isReferenceExpr(item) ? item.target.$refText : undefined));

    const result: Record<string, string> = {};
    referenceNames.forEach((name, i) => {
        if (name) {
            // eslint-disable-next-line @typescript-eslint/no-non-null-assertion
            result[name] = fieldNames[i]!;
        }
    });
    return result;
}

function getDeleteCascades(model: DataModel): string[] {
    const allModels = getDataModels(model.$container);
    return allModels
        .filter((m) => {
            if (m === model) {
                return false;
            }
            const relationFields = m.fields.filter((f) => {
                if (f.type.reference?.ref !== model) {
                    return false;
                }
                const relationAttr = getAttribute(f, '@relation');
                if (relationAttr) {
                    const onDelete = getAttributeArg(relationAttr, 'onDelete');
                    if (onDelete && isEnumFieldReference(onDelete) && onDelete.target.ref?.name === 'Cascade') {
                        return true;
                    }
                }
                return false;
            });
            return relationFields.length > 0;
        })
        .map((m) => m.name);
}

function generateDefaultValueProvider(
    field: DataModelField | TypeDefField,
    extraFunctions: OptionalKind<FunctionDeclarationStructure>[]
) {
    const defaultAttr = getAttribute(field, '@default');
    if (!defaultAttr) {
        return undefined;
    }

    const expr = defaultAttr.args[0]?.value;
    if (!expr) {
        return undefined;
    }

    // find `auth()` in default value expression
    const hasAuth = streamAst(expr).some(isAuthInvocation);
    if (!hasAuth) {
        return undefined;
    }

    // generates a provider function like:
    //     function $default$Model$field(user: any) { ... }
    const funcName = `$default$${field.$container.name}$${field.name}`;
    extraFunctions.push({
        name: funcName,
        parameters: [{ name: 'user', type: 'any' }],
        returnType: 'unknown',
        statements: (writer) => {
            const tsWriter = new TypeScriptExpressionTransformer({ context: ExpressionContext.DefaultValue });
            const code = tsWriter.transform(expr, false);
            writer.write(`return ${code};`);
        },
    });

    return funcName;
}

function isAutoIncrement(field: DataModelField) {
    const defaultAttr = getAttribute(field, '@default');
    if (!defaultAttr) {
        return false;
    }

    const arg = defaultAttr.args[0]?.value;
    if (!arg) {
        return false;
    }

    return isInvocationExpr(arg) && arg.function.$refText === 'autoincrement';
}

function writeShortNameMap(options: ModelMetaGeneratorOptions, writer: CodeWriter) {
    if (options.shortNameMap && options.shortNameMap.size > 0) {
        writer.write('shortNameMap:');
        writer.block(() => {
            for (const [key, value] of options.shortNameMap!) {
                writer.write(`${key}: '${value}',`);
            }
        });
        writer.write(',');
    }
}

function getOnDeleteAction(fieldInfo: DataModelField) {
    const relationAttr = getAttribute(fieldInfo, '@relation');
    if (relationAttr) {
        const onDelete = getAttributeArg(relationAttr, 'onDelete');
        if (onDelete && isEnumFieldReference(onDelete)) {
            return onDelete.target.ref?.name;
        }
    }
    return undefined;
}

function getOnUpdateAction(fieldInfo: DataModelField) {
    const relationAttr = getAttribute(fieldInfo, '@relation');
    if (relationAttr) {
        const onUpdate = getAttributeArg(relationAttr, 'onUpdate');
        if (onUpdate && isEnumFieldReference(onUpdate)) {
            return onUpdate.target.ref?.name;
        }
    }
    return undefined;
}

function exprToValue(value: Expression): unknown {
    return match(value)
        .when(isStringLiteral, (v) => v.value)
        .when(isBooleanLiteral, (v) => v.value)
        .when(isNumberLiteral, (v) => {
            let num = parseInt(v.value);
            if (isNaN(num)) {
                num = parseFloat(v.value);
            }
            if (isNaN(num)) {
                return undefined;
            }
            return num;
        })
        .when(isArrayExpr, (v) => v.items.map((item) => exprToValue(item)))
        .when(isObjectExpr, (v) => exprToObject(v))
        .otherwise(() => undefined);
}

function exprToObject(value: ObjectExpr): unknown {
    return Object.fromEntries(value.fields.map((field) => [field.name, exprToValue(field.value)]));
}
